#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gmssl/sm2.h>
#include <gmssl/pem.h>

#define SM2_CIPHERTEXT_SIZE 1024  // 根据需要调整大小

void print_usage() {
    printf("Usage: sm2_encrypt <public_key.pem> <input_file> <output_file>\n");
}

int main(int argc, char *argv[]) {
    if (argc != 4) {
        print_usage();
        return 1;
    }

    const char *public_key_file = argv[1];
    const char *input_file = argv[2];
    const char *output_file = argv[3];

    // Load public key
    SM2_KEY sm2_key;
    FILE *fp = fopen(public_key_file, "r");
    if (!fp) {
        perror("Failed to open public key file");
        return 1;
    }

    // Use the correct function to read the public key
    if (sm2_public_key_info_from_pem(&sm2_key, fp) != 1) {
        fprintf(stderr, "Failed to read public key from PEM\n");
        fclose(fp);
        return 1;
    }
    fclose(fp);

    // Read input file
    FILE *in_fp = fopen(input_file, "rb");
    if (!in_fp) {
        perror("Failed to open input file");
        return 1;
    }

    fseek(in_fp, 0, SEEK_END);
    long input_len = ftell(in_fp);
    fseek(in_fp, 0, SEEK_SET);
    unsigned char *input_data = malloc(input_len);
    if (fread(input_data, 1, input_len, in_fp) != input_len) {
        perror("Failed to read input file");
        free(input_data);
        fclose(in_fp);
        return 1;
    }
    fclose(in_fp);

    // Encrypt data
    unsigned char ciphertext[SM2_CIPHERTEXT_SIZE];
    size_t ciphertext_len = sizeof(ciphertext);
    
    if (sm2_encrypt(&sm2_key, input_data, input_len, ciphertext, &ciphertext_len) != 1) {
        fprintf(stderr, "SM2 encryption failed\n");
        free(input_data);
        return 1;
    }

    free(input_data);

    // Write output file
    FILE *out_fp = fopen(output_file, "wb");
    if (!out_fp) {
        perror("Failed to open output file");
        return 1;
    }

    if (fwrite(ciphertext, 1, ciphertext_len, out_fp) != ciphertext_len) {
        perror("Failed to write output file");
        fclose(out_fp);
        return 1;
    }
    fclose(out_fp);

    printf("Encryption successful, output written to %s\n", output_file);
    return 0;
}
